import math
import numpy as np
from math import cos, pi, sin
import matplotlib.pyplot as plt


class GetEllipse:

    # modified by Chen Hongjie 2020.08.13
    def get(dataset):
        # fix bug by Chen Hongjie 2020.08.12
        if dataset.shape[0] == 1:
            vec = np.array([dataset[0]])
            dataset = np.append(dataset, vec-0.001, axis=0)
        # if dataset.shape[0] == 2:
        #     meanVec = np.mean(dataset, axis=0)
        #     meanVec = np.array([meanVec])
        #     dataset = np.append(dataset, meanVec+0.001, axis=0)

        dataScale = np.max(dataset, axis=0) - np.min(dataset, axis=0)
        visualScale = np.mean(dataScale) * 0.05

        # calculate the centers of ellipse
        avg_x = sum(element[0] for element in dataset) / len(dataset)
        avg_y = sum(element[1] for element in dataset) / len(dataset)

        # calculate Covariance matrix of the data
        covxy = np.cov(dataset.T)
        # A2 Eigenvalue, B eigenvector
        A2, B = np.linalg.eig(covxy)

        idx_max = np.argmax(A2)
        idx_min = np.argmin(A2)

        A2 = np.maximum(A2, 0)

        # a is major axis of length , b is minor axis of length
        a = math.sqrt(5.991 * A2[idx_max]) + visualScale
        b = math.sqrt(5.991 * A2[idx_min]) + visualScale

        alpha = math.atan(B[1][idx_max] / B[0][idx_max]) ###

        # print('alpha =', alpha)
        listx = []
        listy = []
        for i in range(len(dataset)):
            listx.append(dataset[i][0])
            listy.append(dataset[i][1])

        t_rot = alpha
        t = np.linspace(0, 2 * pi, 100)
        Ell = np.array([a * np.cos(t), b * np.sin(t)])
        R_rot = np.array([[cos(t_rot), -sin(t_rot)], [sin(t_rot), cos(t_rot)]])
        Ell_rot = np.zeros((2, Ell.shape[1]))
        for i in range(Ell.shape[1]):
            Ell_rot[:, i] = np.dot(R_rot, Ell[:, i])

        # ---------- plot
        # plt.plot([0, B[0][idx_max]], [0, B[1][idx_max]], c='r')
        # plt.plot([0, B[0][idx_min]], [0, B[1][idx_min]], c='b')

        # plt.plot(avg_x + Ell_rot[0, :], avg_y + Ell_rot[1, :], 'darkorange')  # rotated ellipse
        # plt.grid(color='lightgray', linestyle='--')
        # plt.scatter(listx, listy, marker="+")
        # plt.axis('scaled')
        # plt.axis('equal')  # changes limits of x or y axis so that equal increments of x and y have the same length
        # plt.show()
        # ------------------------

        data =np.array(list(zip(avg_x + Ell_rot[0, :], avg_y + Ell_rot[1, :])))
        return data


if __name__ == '__main__':
    dataset_test = np.random.rand(50, 2)
    dataset_test.sort(axis=0)
    filterProcess = [
        ['KdeFilter', np.array(range(50))],
        ['DensityPeakFilter', np.array(range(10))],
        ['BayesianNonparemFilter', np.array(range(30, 40))]
    ]
    ellData = [GetEllipse.get(dataset_test[item[1], :]) for item in filterProcess]
    print(ellData)



